import flair
from flair.data import Corpus
from flair.datasets import ColumnCorpus, WNUT_17
from flair.embeddings import TransformerWordEmbeddings
from flair.models import SpanClassifier
from flair.trainers import ModelTrainer

flair.set_seed(123)

corpus: Corpus = WNUT_17().downsample(0.01)

label_dictionary = corpus.make_label_dictionary("ner", add_unk=True)

embeddings = TransformerWordEmbeddings(
    model="distilbert-base-uncased", layers="-1", subtoken_pooling="first", fine_tune=True, use_context=True
)

tagger = SpanClassifier(embeddings=embeddings, label_dictionary=label_dictionary)

trainer = ModelTrainer(tagger, corpus)

trainer.fine_tune(
    "./results/",
    max_epochs=1,
)